{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using TAO YOLOv3\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of YOLOv3\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/sample.jpg?raw=true\" width=\"960\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained ResNet-18 model and train a ResNet-18 YOLO v3 model on the KITTI dataset\n",
    "* Prune the trained YOLO v3 model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the pruned model\n",
    "* Quantize the pruned model using QAT\n",
    "* Run Inference on the trained model\n",
    "* Export the pruned, quantized and retrained model to a .onnx file for deployment to DeepStream\n",
    "* Run inference on the exported .onnx model to verify deployment using TensorRT\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `yolo_v3` model\n",
    "which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of YOLO v3 object detection using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Install the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "     2.1 [Download the dataset](#head-2-1)<br>\n",
    "     2.2 [Verify the downloaded dataset](#head-2-2)<br>\n",
    "     2.4 [Prepare the datasets for training and validation](#head-2-3)<br>\n",
    "     2.3 [Download pretrained model](#head-2-4)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained model](#head-8)\n",
    "9. [Visualize inferences](#head-9)\n",
    "10. [Model Export](#head-10)\n",
    "11. [Verify deployed model](#head-11)\n",
    "12. [QAT workflow](#head-12) <br>\n",
    "    12.1 [QAT Training](#head-12-1) <br>\n",
    "    12.2. [QAT Evaluation](#head-12-2) <br>\n",
    "    12.3. [Pruning QAT model](#head-12-3)<br>\n",
    "    12.4. [Retraining](#head-12-4)<br>\n",
    "    12.5. [Evaluation of the retrained model](#head-12-5)<br>\n",
    "    12.6. [Inference of the retrained QAT model](#head-12-6)<br>\n",
    "    12.7. [Deployment of the QAT model](#head-12-7)<br>\n",
    "    12.8. [Verify the deployed QAT model](#head-12-8)<br>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives<a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/yolo_v3`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/yolo_v3\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/yolo_v3\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/yolo_v3\n",
    "%env LOCAL_PROJECT_DIR=YOUR_LOCAL_PROJECT_DIR_PATH\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"yolo_v3\")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/yolo_v3/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from outside to inside of the docker instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ]\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "After setting up your virtual environment with the above requirements, install TAO pip package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the KITTI detection dataset for the tutorial. To find more details please visit\n",
    " http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download the KITTI detection images (http://www.cvlibs.net/download.php?file=data_object_image_2.zip) and labels (http://www.cvlibs.net/download.php?file=data_object_label_2.zip) to $DATA_DOWNLOAD_DIR.\n",
    " \n",
    " The data will then be extracted to have\n",
    " * training images in `$LOCAL_DATA_DIR/training/image_2`\n",
    " * training labels in `$LOCAL_DATA_DIR/training/label_2`\n",
    " * testing images in `$LOCAL_DATA_DIR/testing/image_2`\n",
    " \n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below.\n",
    "\n",
    "*Note: There are no labels for the testing images, therefore we use it just to visualize inferences for the trained model.*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have gotten the download links in your email, please populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL` and the `KITTI_LABELS_DOWNLOAD_URL`. This next cell, will download the data and place in `$LOCAL_DATA_DIR`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "os.environ[\"URL_IMAGES\"]=KITTI_IMAGES_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $LOCAL_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "os.environ[\"URL_LABELS\"]=KITTI_LABELS_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $LOCAL_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This may take a while: verify integrity of zip files \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_image_2.zip | cut -d ' ' -f 1 | grep -xq '^351c5a2aa0cd9238b50174a3a62b846bc5855da256b82a196431d60ff8d43617$' ; \\\n",
    "if test $? -eq 0; then echo \"images OK\"; else echo \"images corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_image_2.zip; fi \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_label_2.zip | cut -d ' ' -f 1 | grep -xq '^4efc76220d867e1c31bb980bbf8cbc02599f02a9cb4350effa98dbb04aaed880$' ; \\\n",
    "if test $? -eq 0; then echo \"labels OK\"; else echo \"labels corrupt, re-download!\" && rm -f $LOCAL_DATA_DIR/data_object_label_2.zip; fi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_image_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_label_2.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "num_training_images = len(os.listdir(os.path.join(DATA_DIR, \"training/image_2\")))\n",
    "num_training_labels = len(os.listdir(os.path.join(DATA_DIR, \"training/label_2\")))\n",
    "num_testing_images = len(os.listdir(os.path.join(DATA_DIR, \"testing/image_2\")))\n",
    "print(\"Number of images in the train/val set. {}\".format(num_training_images))\n",
    "print(\"Number of labels in the train/val set. {}\".format(num_training_labels))\n",
    "print(\"Number of images in the test set. {}\".format(num_testing_images))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Prepare the datasets for training and validation <a class=\"anchor\" id=\"head-2-3\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Either TFRecord dataset format or the raw KITTI format(with directories of images and KITTI labels) can be used for training/validation. TFRecord dataset format is the recommended one as it is generally faster than reading numerious KITTI label files. However, in some cases, especially when training with small resolutions like 416x416, the experimental result shows TFRecord dataset is almost the same or slightly slower than the other format. So by default, we will use TFRecord dataset in this notebook. As a reference, we still provide the spec files for raw KITTI dataset format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TFrecords conversion spec file for training\")\n",
    "!cat $LOCAL_SPECS_DIR/tfrecords_kitti_trainval.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a new directory for the output tfrecords dump.\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/data\n",
    "# KITTI trainval\n",
    "!tao model yolo_v3 dataset_convert -d $SPECS_DIR/tfrecords_kitti_trainval.txt \\\n",
    "                     -o $DATA_DOWNLOAD_DIR/yolo_v3/kitti_trainval/kitti_trainval \\\n",
    "                     -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, if you have your own dataset already in a volume (or folder), you can mount the volume on `LOCAL_DATA_DIR` (or create a soft link). Below shows an example:\n",
    "```bash\n",
    "# if your dataset is in /dev/sdc1\n",
    "mount /dev/sdc1 $LOCAL_DATA_DIR\n",
    "\n",
    "# if your dataset is in folder /var/dataset\n",
    "ln -sf /var/dataset $LOCAL_DATA_DIR\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you use your own dataset, you will need to run the code below to generate the best anchor shape\n",
    "\n",
    "# !tao model yolo_v3 kmeans -l $DATA_DOWNLOAD_DIR/kitti_split/training/label \\\n",
    "#                     -i $DATA_DOWNLOAD_DIR/kitti_split/training/image \\\n",
    "#                     -n 9 \\\n",
    "#                     -x 1248 \\\n",
    "#                     -y 384\n",
    "\n",
    "# The anchor shape generated by this script is sorted. Write the first 3 into small_anchor_shape in the config\n",
    "# file. Write middle 3 into mid_anchor_shape. Write last 3 into big_anchor_shape."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Download pre-trained model <a class=\"anchor\" id=\"head-2-4\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_object_detection:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_object_detection:resnet18 \\\n",
    "                    --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_object_detection_vresnet18"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Augmentation parameters for on-the-fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.\n",
    "* Whether to use quantization aware training (QAT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Provide pretrained model path\n",
    "!sed -i 's,EXPERIMENT_DIR,'\"$USER_EXPERIMENT_DIR\"',' $LOCAL_SPECS_DIR/yolo_v3_train_resnet18_tfrecord.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/yolo_v3_train_resnet18_tfrecord.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"To run with multigpu, please change --gpus based on the number of available GPUs in your machine.\")\n",
    "!tao model yolo_v3 train -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \\\n",
    "                   --gpus 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"To resume from checkpoint, please change pretrain_model_path to resume_model_path in config file.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/yolov3_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model yolo_v3 evaluate -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/yolov3_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion (`Only for resnets as they have element wise operations or MobileNets.`)\n",
    "* Threshold for pruning.\n",
    "* Output directory to store the model\n",
    "\n",
    "Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold value depends on the dataset and the model. `0.5` in the block below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model yolo_v3 prune -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/yolov3_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                   -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord.txt \\\n",
    "                   -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/yolov3_resnet18_pruned.hdf5 \\\n",
    "                   -eq intersection \\\n",
    "                   -pth 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Printing the retrain spec file. \n",
    "# Here we have updated the spec file to include the newly pruned model as a pretrained weights.\n",
    "!sed -i 's,EXPERIMENT_DIR,'\"$USER_EXPERIMENT_DIR\"',' $LOCAL_SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt\n",
    "!cat $LOCAL_SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model yolo_v3 train --gpus 1 \\\n",
    "                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/yolov3_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 evaluate -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/yolov3_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `infer` tool to generate inferences on the trained models and visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy some test images\n",
    "!mkdir -p $LOCAL_DATA_DIR/test_samples\n",
    "!cp $LOCAL_DATA_DIR/testing/image_2/00000* $LOCAL_DATA_DIR/test_samples/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tao model yolo_v3 inference -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                       -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                       -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/yolov3_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                       -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `inference` tool produces two outputs. \n",
    "1. Overlain images in `$LOCAL_EXPERIMENT_DIR/images_annotated`\n",
    "2. Frame by frame bbox labels in kitti format located in `$LOCAL_EXPERIMENT_DIR/labels`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Model Export <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you trained a non-QAT model, you may export in FP32, FP16 or INT8 mode using the code block below. For INT8, you need to provide calibration image directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# tao <task> export will fail if .onnx already exists. So we clear the export folder before tao <task> export\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/export\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export\n",
    "# Generate .onnx file using tao container\n",
    "!tao model yolo_v3 export -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/yolov3_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                    -o $USER_EXPERIMENT_DIR/export/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                    -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                    --target_opset 12 \\\n",
    "                    --gen_ds_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference.\n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's.\n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP32). \n",
    "!tao deploy yolo_v3 gen_trt_engine -m $USER_EXPERIMENT_DIR/export/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                                   --batch_size 16 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --data_type fp32 \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP16). \n",
    "!tao deploy yolo_v3 gen_trt_engine -m $USER_EXPERIMENT_DIR/export/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                                   --batch_size 16 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --data_type fp16 \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export/trt.engine.fp16 \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To export in INT8 mode (generate calibration cache file). \n",
    "\n",
    "!tao deploy yolo_v3 gen_trt_engine -m $USER_EXPERIMENT_DIR/export/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                                   --cal_image_dir $DATA_DOWNLOAD_DIR/testing/image_2 \\\n",
    "                                   --data_type int8 \\\n",
    "                                   --batch_size 16 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --batches 10 \\\n",
    "                                   --cal_cache_file $USER_EXPERIMENT_DIR/export/cal.bin  \\\n",
    "                                   --cal_data_file $USER_EXPERIMENT_DIR/export/cal.tensorfile \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export/trt.engine.int8 \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Note:` In this example, for ease of execution we restrict the number of calibrating batches to 10. TAO Toolkit recommends the use of at least 10% of the training dataset for int8 calibration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify the deployed model <a class=\"anchor\" id=\"head-11\"></a>\n",
    "Verify the converted engine by visualizing TensorRT inferences.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Infer using TensorRT engine\n",
    "\n",
    "!tao deploy yolo_v3 inference -m $USER_EXPERIMENT_DIR/export/trt.engine \\\n",
    "                              -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord.txt \\\n",
    "                              -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                              -r $USER_EXPERIMENT_DIR/yolo_infer_images \\\n",
    "                              -t 0.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'yolo_infer_images/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. QAT workflow <a class=\"anchor\" id=\"head-12\"></a>\n",
    "In this section, we will explore the typical Quantization-Aware Training(QAT) workflow with TAO. QAT workflow is almost the same as non-QAT workflow except for two major differences:\n",
    "1. set `enable_qat` to `True` in training and retraining spec files to enable the QAT for training/retraining\n",
    "2. when doing export in INT8 mode, the calibration json file that stores the scales used during QAT is extracted during `tao <task> export`. The .onnx file and calibration json file are used to generate engine file through tao deploy"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.1. QAT Training <a class=\"anchor\" id=\"head-12-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To enable QAT training on sample spec file, uncomment following lines\n",
    "!sed -i \"s/enable_qat: false/enable_qat: true/g\" $LOCAL_SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt\n",
    "!sed -i 's,EXPERIMENT_DIR,'\"$USER_EXPERIMENT_DIR\"',' $LOCAL_SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt\n",
    "!cat $LOCAL_SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned_QAT\n",
    "print(\"To run with multigpu, please change --gpus based on the number of available GPUs in your machine.\")\n",
    "!tao model yolo_v3 train -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned_QAT \\\n",
    "                   --gpus 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned_QAT/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now check the evaluation stats in the csv file and pick the model with highest eval accuracy.\n",
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned_QAT/yolov3_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.2. QAT Evaluation <a class=\"anchor\" id=\"head-12-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 evaluate -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned_QAT/weights/yolov3_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.3. Pruning QAT model <a class=\"anchor\" id=\"head-12-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned_QAT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 prune -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned_QAT/weights/yolov3_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                   -e $SPECS_DIR/yolo_v3_train_resnet18_tfrecord_qat.txt \\\n",
    "                   -o $USER_EXPERIMENT_DIR/experiment_dir_pruned_QAT/yolov3_resnet18_pruned.hdf5 \\\n",
    "                   -eq intersection \\\n",
    "                   -pth 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned_QAT/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.4. Retraining <a class=\"anchor\" id=\"head-12-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -i \"s/enable_qat: false/enable_qat: true/g\" $LOCAL_SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt\n",
    "!sed -i 's,EXPERIMENT_DIR,'\"$USER_EXPERIMENT_DIR\"',' $LOCAL_SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model yolo_v3 train --gpus 1 \\\n",
    "                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                   -r $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat/yolov3_training_log_resnet18.csv\n",
    "%set_env EPOCH=080"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.5. Evaluation of the retrained model <a class=\"anchor\" id=\"head-12-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 evaluate -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                      -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/yolov3_resnet18_epoch_$EPOCH.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.6. Inference of the retrained QAT model <a class=\"anchor\" id=\"head-12-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 inference -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                       -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                       -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/yolov3_resnet18_epoch_$EPOCH.hdf5 \\\n",
    "                       -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_PATH = 'images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.7. Deployment of the QAT model <a class=\"anchor\" id=\"head-12-7\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " #### Generate .onnx file using tao container\n",
    "If you train a QAT model, you may only export in INT8 mode using following code block. This generates an onnx file and the corresponding calibration json file that stores scales used during QAT. You can either use the onnx file and calibration json file to generate int8 engine through tao deploy or DeepStream for FP32 or FP16 mode. But please note this gives sub-optimal results. If you want to deploy in FP32 or FP16, you should disable QAT in training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export_qat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model yolo_v3 export -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/yolov3_resnet18_epoch_$EPOCH.hdf5  \\\n",
    "                    -o $USER_EXPERIMENT_DIR/export_qat/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                    -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                    --cal_json_file $USER_EXPERIMENT_DIR/export_qat/cal.json \\\n",
    "                    --target_opset 12 \\\n",
    "                    --gen_ds_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export_qat"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " #### Generate a TensorRT engine using tao deploy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy yolo_v3 gen_trt_engine -m $USER_EXPERIMENT_DIR/export_qat/yolov3_resnet18_epoch_$EPOCH.onnx \\\n",
    "                                   -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                                   --data_type int8 \\\n",
    "                                   --batch_size 8 \\\n",
    "                                   --min_batch_size 1 \\\n",
    "                                   --opt_batch_size 8 \\\n",
    "                                   --max_batch_size 16 \\\n",
    "                                   --cal_cache_file $USER_EXPERIMENT_DIR/export_qat/cal.bin \\\n",
    "                                   --cal_json_file $USER_EXPERIMENT_DIR/export_qat/cal.json \\\n",
    "                                   --engine_file $USER_EXPERIMENT_DIR/export_qat/trt.engine.int8 \\\n",
    "                                   --results_dir $USER_EXPERIMENT_DIR/export_qat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported engine:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/export_qat/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### 12.8. Verify the deployed QAT model <a class=\"anchor\" id=\"head-12-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy yolo_v3 inference -m $USER_EXPERIMENT_DIR/export_qat/trt.engine.int8 \\\n",
    "                              -e $SPECS_DIR/yolo_v3_retrain_resnet18_tfrecord_qat.txt \\\n",
    "                              -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                              -r $USER_EXPERIMENT_DIR/yolo_infer_images_qat_trt \\\n",
    "                              -t 0.6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'yolo_infer_images_qat_trt/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 3 # number of columns in the visualizer grid.\n",
    "IMAGES = 9 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
